# -*- coding = utf-8 -*-
# @Time : 2022/4/12 9:52
# @Author : Yuan
# @File : recognize.py
# @Software : PyCharm

import sys
# from keras.engine import input_spec
import numpy as np
from matplotlib import pyplot
# from keras.utils import to_categorical

from keras.models import Sequential
from keras.layers import Conv2D
from keras.layers import MaxPooling2D
from keras.layers import Dense
from keras.layers import Flatten
from tensorflow.keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
import matplotlib
from matplotlib.image import pil_to_array

import tensorflow as tf
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)


# 创建一个cnn模型：
def define_cnn_model():
    # 使用序列模型
    model = Sequential()
    # 卷积层
    model.add(Conv2D(32, (3, 3), activation='relu', padding='same', input_shape=(200, 200, 3)))
    # 最大池化层
    model.add(MaxPooling2D((2, 2)))
    # Flatten层
    model.add(Flatten())
    # 全连接层
    model.add(Dense(128, activation="relu"))
    model.add(Dense(1, activation="sigmoid"))

    # 编译模型
    opt = SGD(lr=0.001, momentum=0.9)
    model.compile(optimizer=opt, loss='binary_crossentropy', metrics=['accuracy'])

    return model



# 实例化训练代码
def train_cnn_model():
    # 实例化模型
    model = define_cnn_model()

    # 创建图片生成器
    datagen = ImageDataGenerator(rescale=1.0 / 255.0)  # 不断产生图片生成

    train_it = datagen.flow_from_directory(
        "./kaggle/classes2",
        class_mode='binary',
        batch_size=64,  # 每次训练64图
        target_size=(200, 200)  # 先把图片进行缩放为200*200的
    )

    # 训练模型
    model.fit_generator(
        train_it,  # 刚才定义的图片生成器
        steps_per_epoch=len(train_it),
        epochs=1,  # 训练轮次
        verbose=1
    )

    return model


# 训练模型
from keras.models import save_model
# save_model(train_cnn_model(), 'my_model.h5')

# 完成训练后，载入刚才训练后保存的模型
from keras.models import load_model
# 修改model_path为你自己保存的模型的位置
model = load_model("./my_model.h5")



# 然后定义一个函数，从测试文件夹中读取任意一张图片。
import os, random
from matplotlib.pyplot import imshow
import numpy as np
from PIL import Image
# %matplotlib inline


def read_random_image():
    folder = r"./kaggle/classes1/test/"
    file_path = folder + random.choice(os.listdir(folder))
    print(file_path)
    pil_im = Image.open(file_path, 'r')
    return pil_im


# 定一个使用模型对读取的图片进行预测的函数
def get_predict(pil_im, model):
    # 对图片进行缩放
    pil_im = pil_im.resize((200, 200))
    # 将格式转为 numpy array 格式
    array_im = np.asarray(pil_im)

    array_im = np.expand_dims(array_im, axis=0)
    # print(array_im/255)
    # 对图片进行预测
    result = model.predict(array_im)
    print(result[0][0])
    if result[0][0] > 0.5:
        print("预测结果是：狗")
    else:
        print("预测结果是：猫")


# -------------------------------------------------------------------------------------


# # # 完成上述代码后，就可以进行代码实测
pil_im = read_random_image()
imshow(np.asarray(pil_im))
get_predict(pil_im, model)



